# -*- coding: utf-8 -*-
# created on 2016/11/30
import copy

from mathsolver.functions.shulie.base import ShuLieNXiangsEqsTransformation
from mathsolver.functions.shulie.basic import *
from mathsolver.functions.mathematica.mathematicaSolve import MathematicaSolve
from mathsolver.functions.shulie.ShuLie_Transform import *


class ShuLieGetVariable001(BaseFunction):
    """
    eg. 设S_{n}为等差数列{a_{n}}的前n项和,若a_{1}=1,公差d=2,且S_{k+2}-S_{k}=24,则k=()
    1. 求变量的值
    2. 单数列
    3. 等差，等比数列类型
    """

    def solver(self, *args):
        assert len(args) == 1
        known = copy.deepcopy(self.known)
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        new_seqs, new_other_seqs = sl_seqs_other_seqs(self.known, seqs)
        assert not new_other_seqs
        assert new_seqs
        if known['inequations']:
            for ineq in known['inequations']:
                left, op, right = ineq
                new_ineq = [sympify(left), op, sympify(right)]
                if new_ineq not in new_seqs:
                    new_seqs.append(new_ineq)
        if known['equations']:
            for eq in known['equations']:
                if eq not in new_seqs:
                    new_seqs.append(eq)
        new_targets = []
        for target in args[0]:
            new_targets.append(target.sympify())
        sls, sl_flag = sl_shulies_flag(self.known, new_seqs)
        assert sum(sl_flag) == 0
        sls = [self.known[item] for item in self.known if isinstance(self.known[item], BaseSequence)]
        stepsolver = DengChaBiJiChuAssume(known=self.known).solver(sls)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        stepsolver = DengChaBiEqNormalization(known=self.known).solver(new_seqs)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        stepsolver = MathematicaSolve().solver(stepsolver.output[0], BaseVariables(new_targets))
        if isinstance(stepsolver.output[0], BaseEq) == 1:
            self.steps.append(["", "解得: %s" % (stepsolver.output[0].printing())])
            self.label.add('解方程')
        else:
            self.steps.append(["", "解得: %s" % (stepsolver.output[0].printing())])
            self.label.add('解方程组')
        solutions = stepsolver.output[0]
        answers = {}
        for new_target in new_targets:
            new_target_values = sl_solutuons_subs(solutions, new_target)
            answers[new_target] = new_target_values
        self.output.append(BaseSymbolValue(answers))
        return self


class ShuLieGetVariable002(BaseFunction):
    """
    在数列{a_{n}}中,a_{1}=2,a_{n+1}=2a_{n},S_{n}为{a_{n}}的前n项和,若S_{n}=126,则n=().
    1. 求变量的值
    2. 单数列
    3. 一般数列
    4. 包含nXiangs_Eqs
    """
    def solver(self, *args):
        assert len(args) == 1
        known = copy.deepcopy(self.known)
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        new_seqs, new_other_seqs = sl_seqs_other_seqs(known, seqs)
        assert not new_other_seqs
        assert new_seqs
        new_targets = []
        for target in args[0]:
            new_targets.append(target.sympify())
        sls, sl_flag = sl_shulies_flag(known, new_seqs)
        assert sum(sl_flag) > 0
        assert len(sls) == 1
        sl = sls[0]
        sl.eqs_update(new_seqs)
        sl.get_shouxiang()
        sl.get_xiangs()
        sl.get_nxiangs()
        sl.get_dituigongshi()
        sl.get_nineqs()
        assert sl.nXiangs_Eqs
        n_xiangs_eqs = sl.nXiangs_Eqs
        sl.get_init()
        stepsolver = ShuLieNXiangsEqsTransformation(known).solver(n_xiangs_eqs)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        condition_eqs = stepsolver.output[0].sympify()
        self.steps.append(["", "∴ %s" % (BaseIneqs(condition_eqs).printing())])
        try:
            solutions = MathematicaSolve().solver(BaseIneqs(condition_eqs)).output[0]
            assert solutions.sympify()
            self.steps.append(["", "解得: %s" % (BaseSymbolValue(solutions.sympify()).printing())])
            answers = {}
            for new_target in new_targets:
                new_target_flags = []
                new_target_values = sl_solutuons_subs(solutions, new_target)
                for new_target_value in new_target_values:
                    if str(new_target_value) == str(new_target):
                        new_target_flags.append(1)
                    else:
                        new_target_flags.append(0)
                assert sum(new_target_flags) == 0
                answers[new_target] = new_target_values
                self.steps.append(
                    ["", "∴%s" % "或".join(["%s = %s" % (new_target, new_latex(v)) for v in new_target_values])])
            self.output.append(BaseSymbolValue(answers))
        except Exception:
            assert len(new_targets) == 1
            target = new_targets[0]
            assert str(target) == str("n")
            assert len(condition_eqs) == 1
            condition_eq = condition_eqs[0]
            answer = None
            for i in range(1, 10000):
                left = condition_eq[0]
                right = condition_eq[1]
                new_left = left.subs({target: i})
                new_right = right.subs({target: i})
                if new_left == new_right:
                    answer = i
                    break
            self.steps.append(["", "解得: %s = %s" % (new_latex(target), new_latex(answer))])
            self.output.append(BaseSymbolValue({target: answer}))
        return self


class ShuLieGetVariable003(BaseFunction):
    """
    已知数列{a_{n}}的前n项和S_{n}=1+λa_{n},其中λ≠0.若S_{5}=\\frac{31}{32},求λ.
    1. 求变量的值
    2. 单数列
    3. 一般数列
    4. 包含Xiangs_Eqs
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        new_seqs, new_other_seqs = sl_seqs_other_seqs(known, seqs)
        assert not new_other_seqs
        assert new_seqs
        new_targets = []
        for target in args[0]:
            new_targets.append(target.sympify())
        sls, sl_flag = sl_shulies_flag(known, new_seqs)
        assert sum(sl_flag) > 0
        assert len(sls) == 1
        sl = sls[0]
        sl.eqs_update(new_seqs)
        sl.get_shouxiang()
        sl.get_nxiangs()
        sl.get_dituigongshi()
        sl.get_xiangs()
        sl.get_nineqs()
        assert sl.Xiangs_Eqs
        xiangs_eqs = sl.Xiangs_Eqs
        new_known = dict(known)
        seqs = []
        for eq in new_known['Sequations']:
            if eq != xiangs_eqs[0]:
                seqs.append(eq)
        new_known['Sequations'] = seqs
        sl.Eqs = []
        sl.a1 = []
        sl.DiTuiGongShi_Eqs = []
        sl.Xiangs_Eqs = []
        sl.DiTuiGongShi_Ineqs = []
        sl.nXiangs_Eqs = []
        stepsolver = ShuLieNXiangsEqsTransformation(new_known).solver(xiangs_eqs)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        condition_eqs = stepsolver.output[0].sympify()
        self.steps.append(["", "∴ %s" % (BaseIneqs(condition_eqs).printing())])
        try:
            solutions = MathematicaSolve().solver(BaseIneqs(condition_eqs)).output[0]
            assert solutions.sympify()
            self.steps.append(["", "解得: %s" % (BaseSymbolValue(solutions.sympify()).printing())])
            answers = {}
            for new_target in new_targets:
                new_target_flags = []
                new_target_values = sl_solutuons_subs(solutions, new_target)
                for new_target_value in new_target_values:
                    if str(new_target_value) == str(new_target):
                        new_target_flags.append(1)
                    else:
                        new_target_flags.append(0)
                assert sum(new_target_flags) == 0
                answers[new_target] = new_target_values
                self.steps.append(["", "∴%s" % "或".join(
                    ["%s = %s" % (new_latex(new_target), new_latex(v)) for v in new_target_values])])
            self.output.append(BaseSymbolValue(answers))
        except Exception:
            assert str(target) == str("n")
            assert len(condition_eqs) == 1
            condition_eq = condition_eqs[0]
            answer = None
            for i in range(1, 10000):
                left = condition_eq[0]
                right = condition_eq[1]
                new_left = left.subs({target: i})
                new_right = right.subs({target: i})
                if new_left == new_right:
                    answer = i
                    break
            self.steps.append(["", "解得: %s = %s" % (new_latex(target), new_latex(answer))])
            self.output.append(BaseSymbolValue({target: answer}))
        return self


class ShuLieGetVariable(BaseFunction):
    CLS = [ShuLieGetVariable001, ShuLieGetVariable002, ShuLieGetVariable003]

    def solver(self, *args):
        known = self.known
        r = None
        for cl in ShuLieGetVariable.CLS:
            try:
                new_known = dict(known)
                r = cl(new_known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r


class ShuLieGetVariables(BaseFunction):
    def solver(self, *args):
        known = self.known
        polys = args[0].sympify()
        for poly in polys:
            new_known = dict(known)
            stepsolver = ShuLieGetVariable(new_known).solver(BasePoly(poly))
            self.steps += stepsolver.steps
        return self


class ShuLieGetGongBi(BaseFunction):
    def solver(self, *args):
        seq = args[0].sympify()
        seq_name = seq.func
        ShuLie = self.search(str(seq_name))
        ShuLieType = ShuLie.type
        if ShuLieType == "DengBi":
            if 'Sequations' in self.known:
                if self.known['Sequations']:
                    Seqs = self.known['Sequations']
                    condition_eqs = []
                    for eq in Seqs:
                        if len(eq) == 2:
                            stepsolver = DengChaBiEqNormalization(self.known).solver(BaseEq(eq))
                            self.label.update(stepsolver.label)
                            new_eq = stepsolver.output[0].sympify()
                            condition_eqs.extend(new_eq)
                        elif len(eq) == 3:
                            stepsolver = DengChaBiEqNormalization(self.known).solver(BaseIneq(eq))
                            self.label.update(stepsolver.label)
                            new_eq = stepsolver.output[0].sympify()
                            condition_eqs.extend(new_eq)
                    self.steps.append(["", "依题意，得"])
                    if len(condition_eqs) == 1:
                        self.steps.append(["", "%s" % (BaseEq(condition_eqs[0]).printing())])
                        stepsolver = MathematicaSolve().solver(BaseEq(condition_eqs[0]))
                    elif len(condition_eqs) > 1:
                        self.steps.append(["", "%s" % (BaseIneqs(condition_eqs).printing())])
                        stepsolver = MathematicaSolve().solver(BaseIneqs(condition_eqs))
                    solutions = stepsolver.output[0]
                    self.steps.append(["", "解得: %s" % (solutions.printing())])
                    qName = ShuLie.qName
                    answer = sl_search_multiple(solutions.sympify(), qName)
                    q_values = []
                    for ans in answer:
                        new_value = qName
                        new_value = new_value.subs(ans)
                        q_values.append(new_value)
                    q_values = set(q_values)
                    q_values = list(q_values)
                    self.steps.append(
                        ["", "∴%s" % "或".join(["%s = %s" % (new_latex(qName), new_latex(v)) for v in q_values])])
        elif ShuLieType == "ShuLie":
            pass
        return self
